package qmlt.learning.decisiontree;

import java.util.Collection;

import qmlt.dataset.Attribute;
import qmlt.dataset.DataSet;
import qmlt.dataset.Instance;
import qmlt.learning.Learner;
import qmlt.learning.decisiontree.control.DecisionTreeController;

public class DecisionTree implements Learner<DecisionTreeController>
{

    private boolean pruningPhase = false;
    private DecisionTreeController controller;
    private Node root;
    
    public Node getRoot()
    {
    	return root;
    }
    
    @Override
    public Object predict(Instance instance)
    {
        return predictHelper(root, instance);
    }
    
    public Object predictHelper(Node node, Instance instance)
    {
        // keep track of predicted instances that go through this node for pruning
        if (pruningPhase)
            node.predictedInstances.add(instance);

        if (node.isLeaf())
        {
            return node.targetClass;
        }

        int i = node.splittingAttributeIndex;
        String v = (String) instance.getFeatures().get(i);
        if (node.children.containsKey(v))
        {
            Node child = node.children.get(v);
            return predictHelper(child, instance);
        }
        else
        {
            System.err.format("warning: unseen attribute value encountered in predicting: %d:%s\n",
                    i, v);
            Node child = getMaxSubsetChild(node.children.values());
            return predictHelper(child, instance);
        }
    }

    protected Node getMaxSubsetChild(Collection<Node> children)
    {
        Node rst = null;
        int maxSize = -1;

        for (Node child : children)
        {
            if (child.instances.size() > maxSize)
            {
                rst = child;
                maxSize = child.instances.size();
            }
        }

        return rst;
    }

    @Override
    public void train(DataSet trainSet, DecisionTreeController controller)
    {
        assert (trainSet.getTargetDef().equals(Attribute.STRING)) :
            "decision tree applies only to classification problems.";
        
        this.controller = controller;
        root = new Node(trainSet.getInstances());
        buildHelper(root);
        
        pruningPhase = true;
        controller.prune(this);
        pruningPhase = false;
    }

    protected void buildHelper(Node node)
    {
        controller.split(node);
        for (Node child : node.children.values())
        {
            buildHelper(child);
        }
    }

		@Override
		public void loadModel(String inputFilepath)
		{
			// TODO Auto-generated method stub
			
		}

		@Override
		public void saveModel(String outputFilepath)
		{
			// TODO Auto-generated method stub
			
		}

}
